/*
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package net.javlov.vf.fa;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import net.javlov.Option;
import net.javlov.SimpleAgent;
import net.javlov.State;
import net.javlov.vf.IneligibleOptionException;
import net.javlov.vf.OptionValues;
import net.javlov.vf.QFunction;
import net.javlov.vf.TableOptionValues;
import net.javlov.vf.UnknownStateException;
import net.javlov.vf.Value;

public class TileCodedQ implements QFunction<double[]>{
	
	private static final long serialVersionUID = 1598954347408177621L;
	
	/**
	 * The TileCoder provides a default generalization of 1 in each state dimension.
	 * The scaler serves to transform the state space for custom generalization. Of
	 * course, it's faster if the whole state space is transformed a priori; then no
	 * scaling will have to be done on each call to the tilecoder.
	 */
	protected Function scaler;
    
	/**
	 * The TileCoder that maps states to (hashed) tile indices.
	 */
    protected TileCoder tiles;
    
    /**
     * Number of overlapping tilings and total number of tiles, for the TileCoder.
     */
    protected int numTilings, memSize;
    
    /**
     * Function weights.
     */
    protected double[][] weights;
        
    protected List<? extends Option<?>> optionPool;
    
    /**
     * Used for storing tile indices, overwritten each time the function is
     * queried / updated.
     */
    private int[] tiling;
        
    public TileCodedQ(int numTilings, int memSize, List<? extends Option<?>> optionPool) {
    	this( new TileCoder(numTilings, memSize), numTilings, memSize, optionPool );    	
    }
    
    public TileCodedQ(TileCoder tc, int numTilings, int memSize, List<? extends Option<?>> optionPool) {
    	tiles = tc;
    	this.numTilings = numTilings;
    	this.memSize = memSize;
    	weights = new double[optionPool.size()][memSize];
    	tiling = new int[numTilings];
    	this.optionPool = optionPool;
    }
    
    public TileCodedQ(int numTilings, int memSize, double[] scaling, List<? extends Option<?>> optionPool) {
    	this(numTilings, memSize, optionPool);
    	setScaler( new ScaleFunction(scaling) );
    }
    
    public TileCodedQ(int numTilings, int memSize, Function scaler, List<? extends Option<?>> optionPool) {
    	this(numTilings, memSize, optionPool);
    	setScaler(scaler);
    }
    
    public void setScaler(Function sc) {
    	scaler = sc;
    }
    
    @Override
	public double getValue(State<? extends double[]> s, Option<?> o) {
    	try {
			return getValues(s).get(o);
		} catch ( IneligibleOptionException e ) {
			throw new IneligibleOptionException(o + " not eligible in " + s);
		}

	}

	@Override
	public OptionValues getValues(State<? extends double[]> s) {
		getTileCode(s);
		return new TCOptionValues( s, optionPool, tiling );
	}

	@Override
	public Double queryValue(State<? extends double[]> s, Option<?> o) {
		return getValue(s,o);
	}

	@Override
	public void init() {
		weights = new double[optionPool.size()][memSize];
	}

	@Override
	public void reset() {
		//do nothing
	}

	@Override
	public void setValue(State<? extends double[]> s, Option<?> o, double value) {
		throw new UnsupportedOperationException();		
	}

	@Override
	public void update(State<? extends double[]> s, Option<?> o, double delta) {
		getTileCode(s);
		new TCOptionValues( s, optionPool, tiling ).update(o, delta);
	}
    
	protected void getTileCode(State<? extends double[]> s) {
		if ( scaler == null )
			tiling = tiles.getTiles(s.getData());
		else
			tiling = tiles.getTiles(scaler.eval( s.getData() ));
	}
    
	
    static class ScaleFunction implements Function {

    	protected double[] sc, ret;
    	
    	public ScaleFunction(double[] scaling) {
    		sc = scaling;
    		ret = new double[scaling.length];
    	}
    	
		@Override
		public double[] eval(double[] input) {
			for ( int i = 0; i < input.length; i++ )
				ret[i] = input[i]*sc[i];
			return ret;
		}
    	
    }
   
    
    class TCOptionValues implements Serializable, OptionValues {

    	private static final long serialVersionUID = -175768075273183858L;
    	
    	private int persistentTiling[];
    	protected List<Option<?>> allOpts;
    	
    	public TCOptionValues(State s, List<? extends Option> optionPool, int[] tiling) {
    		persistentTiling = Arrays.copyOf(tiling, tiling.length);
    		allOpts = s.getOptionSet();
/*    		allOpts = new ArrayList<Option>();
    		
    		for ( Option o : optionPool )
    			if ( o.isEligible(s) )
    				addOption(s, o);*/
    		
    		if ( allOpts.isEmpty() )
    			throw new IllegalStateException("No eligible options for state " + s);
    	}

    	
    	@Override
		public double get(Option o) {	
			double sum = 0;
			for ( int i = 0; i < numTilings; i++ )
				sum += weights[o.getID()][persistentTiling[i]];
			return sum;
		}
    	    
    	@Override
		public void update(Option o, double delta) {
			for ( int i = 0; i < numTilings; i++ )
				weights[o.getID()][persistentTiling[i]] += delta/numTilings;
		}
    	    	
    	/* (non-Javadoc)
    	 * @see net.javlov.vf.OptionValues#getMaxOptions()
    	 */
    	@Override
    	public List<? extends Option<?>> getMaxOptions() {
    		double sum, maxValue = Double.NEGATIVE_INFINITY;
			List<Option<?>> maxOpts = new ArrayList<Option<?>>(allOpts.size());
			for ( Option o : allOpts ) {
				sum = get(o);
				if ( sum > maxValue ) {
    				maxOpts.clear();
    				maxOpts.add(o);
    				maxValue = sum;
    			} else if ( sum == maxValue )
    				maxOpts.add(o);
			}
			return maxOpts;
    	}
    	
    	/* (non-Javadoc)
    	 * @see net.javlov.vf.OptionValues#getOptions()
    	 */
    	@Override
    	public List<? extends Option<?>> getOptions() {
    		return allOpts;
    	}
    	
    	/* (non-Javadoc)
    	 * @see net.javlov.vf.OptionValues#getValues()
    	 */
    	@Override
    	public Collection<Value> getValues() {
    		List<Value> vals = new ArrayList<Value>(allOpts.size());
    		double sum;
    		for ( Option o : allOpts ) {
				sum = get(o);
				vals.add( new Value(sum) );
    		}
    		return vals;
    	}
    	
    	/* (non-Javadoc)
    	 * @see net.javlov.vf.OptionValues#size()
    	 */
    	@Override
    	public int size() {
    		return allOpts.size();
    	}
    	
    	/* (non-Javadoc)
    	 * @see net.javlov.vf.OptionValues#max()
    	 */
    	@Override
    	public double max() {
    		double sum,
    				maxValue = Double.NEGATIVE_INFINITY;
			for ( Option o : allOpts ) {
				sum = get(o);
				if ( sum > maxValue )
    				maxValue = sum;
			}
			return maxValue;
    	}
    	
    	public String toString() {
    		return getValues().toString();
    	}

    }
}
